import fitz  # PyMuPDF
import sys
import os

pdf_path = sys.argv[1]
output_path = sys.argv[2]
fileHash = sys.argv[3]
convert_type = sys.argv[4]
output_folder = f"{output_path}/{fileHash}"

# 用于提取PDF中的所有图片
def extract_images_from_pdf(pdf_path, output_folder):
    # 打开PDF文件
    pdf = fitz.open(pdf_path)
    image_list = []
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    # 遍历PDF的每一页
    for page_number in range(len(pdf)):
        # 获取当前页
        page = pdf.load_page(page_number)
        
        # 获取页面中的图片
        image_list = page.get_images(full=True)
        
        # 遍历所有图片
        for image_index, img in enumerate(image_list, start=1):
            # 获取图片信息
            xref = img[0]
            base_image = pdf.extract_image(xref)
            image_bytes = base_image["image"]
            
            # 获取图片的扩展名
            ext = base_image["ext"]
            image_filename = f"{output_folder}/image_{page_number+1}_{image_index}.{ext}"
            
            # 保存图片
            with open(image_filename, "wb") as image_file:
                image_file.write(image_bytes)
    
    # 关闭PDF文件
    pdf.close()
    

# 用于将PDF转换为图片
def pdf_to_images(pdf_path, output_folder):
    # 打开PDF文件
    pdf = fitz.open(pdf_path)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    for page_number in range(len(pdf)):
        # 获取PDF页面
        page = pdf[page_number]
        # 渲染页面为图片
        pix = page.get_pixmap()
        # 保存图片
        image_path = f"{output_folder}/{fileHash}_page_{page_number + 1}.png"
        pix.save(image_path)
    pdf.close()

if convert_type == "pdf_image":
    extract_images_from_pdf(pdf_path, output_folder)
elif convert_type == "page_image":
    pdf_to_images(pdf_path, output_folder)